from easy_io import read_pkl_file, write_pkl_file


def change_fold(nodule_list, target_fold):
    for inodule in range(len(nodule_list)):
        if nodule_list[inodule]['scanid'] in target_fold:
            nodule_list[inodule]['fold'] = target_fold[nodule_list[inodule]['scanid']]['fold']
            nodule_list[inodule]['cancer'] = target_fold[nodule_list[inodule]['scanid']]['cancer']
        else:
            nodule_list[inodule]['fold'] = -1


def main(cdd_file_list, out_file_name):
    foldlist = read_pkl_file('files/kaggle_training_set_4_fold_list.pkl')
    all_candidates_for_training_mal = []
    for cdd_file in cdd_file_list:
        cdds = read_pkl_file(cdd_file)
        change_fold(cdds, foldlist)
        all_candidates_for_training_mal.extend(list(cdds))
    write_pkl_file(out_file_name, all_candidates_for_training_mal)

    # for x in all_candidates_for_training_mal:
    #     if 'cancer' in x:
    #         print(x['scanid'], x['fold'], x['cancer'])
    #     else:
    #         print(x['scanid'], x['fold'])


if __name__ == '__main__':
    cdd_file_list = ['../../data/reduced_candidates/3dcnn_kaggle_testset_candidates_v4_newlabel.pkl',
                     '../../data/reduced_candidates/3dcnn_candidates_list_v9_newlabel.pkl',
                     '../../data/reduced_candidates/3dcnn_kaggle_beni_candidates_v4.pkl',
                     '../../data/reduced_candidates/3dcnn_spie_candidates_v4_newlabel.pkl']
    out_file_name = '../../data/tmp_classification_branch1/data/3dcnn_candidates_list_v9_4fold_newfold_newlabel.pkl'
    main(cdd_file_list, out_file_name)

    cdd_file_list = ['../../data/reduced_candidates/3dcnn_cddv15_v2_kaggle_beni_candidates_v10.pkl',
                     '../../data/reduced_candidates/3dcnn_cddv15_v2_kaggle_testset_candidates_v10.pkl',
                     '../../data/reduced_candidates/3dcnn_cddv15_v2_lidc_kaggle_candidates_v15.pkl',
                     '../../data/reduced_candidates/3dcnn_cddv15_v2_spie_candidates_v10.pkl']
    out_file_name = '../../data/tmp_classification_branch1/data/3dcnn_cddv15_v2_4fold_alldata.pkl'
    main(cdd_file_list, out_file_name)